"""
将 MIDOG 2021 中  tif 转换成 yolo 系列需要的数据集

https://midog.deepmicroscopy.org/download-dataset/

"""
import json
import os
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm

if __name__ == '__main__':
    midog = 'MIDOG2021'
    wsi_path = f"/media/hsmy/wanghao_18T/dataset/{midog}/tiff/"
    png_path = f"/media/hsmy/wanghao_18T/dataset/{midog}/yolo_20_256_t/images/"
    label_path = f"/media/hsmy/wanghao_18T/dataset/{midog}/yolo_20_256_t/labels/"
    os.makedirs(png_path, exist_ok=True)
    os.makedirs(label_path, exist_ok=True)
    midog_json = f"/media/hsmy/wanghao_18T/dataset/{midog}/MIDOG.json"
    data = json.load(open(midog_json))

    # mask_path = '/media/hsmy/wanghao_18T/dataset/MIDOG2022/yolo_10_256_origin/masks/'
    # os.makedirs(mask_path, exist_ok=True)

    patch_size = 256  # 训练 image 大小
    react_size = 64  # 框的大小 40X 下是64x64
    level_10X = False  # level
    level_20X = True  # level
    if level_10X:
        react_size = react_size / 4  # 框的大小 10X 下是64x64 / 4
    if level_20X:
        react_size = react_size / 2  # 框的大小 10X 下是64x64 / 4
    react_size_percent = react_size / patch_size

    images_dict = {image["file_name"]: image["id"] for image in data["images"]}
    for wsi_file in tqdm(os.listdir(wsi_path)):
        image_id = images_dict[wsi_file]
        # test_arr = [6, 7, 14, 15, 18, 26, 29, 31, 32, 35, 50, 52, 57, 58, 59, 64, 66, 74, 88, 93, 95, 100, 102, 104,
        #             107, 109, 115, 116, 123, 124, 128, 147, 150]
        # if image_id in test_arr:
        #     continue
        file_name = wsi_file.split('.')[0]
        file_path = os.path.join(wsi_path, wsi_file)
        img = cv2.imread(file_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        height, width, _ = img.shape

        if level_10X:
            img = cv2.resize(img, (width // 4, height // 4), interpolation=cv2.INTER_NEAREST)
            height, width, _ = img.shape

        if level_20X:
            img = cv2.resize(img, (width // 2, height // 2), interpolation=cv2.INTER_NEAREST)
            height, width, _ = img.shape

        mask = np.zeros((height, width)).astype(np.uint8)

        annotations = [anno for anno in data["annotations"] if
                       anno["image_id"] == image_id and anno["category_id"] == 1]

        if len(annotations) == 0:
            continue

        for index, anno in enumerate(annotations):
            bbox_arr = anno["bbox"]
            # 二维数组的x,y是反过来的
            if level_10X:
                x = bbox_arr[1] // 4
                y = bbox_arr[0] // 4

                x = int(x + 5)
                y = int(y + 5)
            elif level_20X:
                x = bbox_arr[1] // 2
                y = bbox_arr[0] // 2

                x = int(x + 10)
                y = int(y + 10)
            else:
                x = bbox_arr[1]
                y = bbox_arr[0]

                x = int(x + 25)
                y = int(y + 25)

            try:
                mask[x, y] = 1
            except Exception as e:
                print(e)
                print(mask.shape)
                print(wsi_file)
                print(f"{x}_{y}")

        index = 0  # 平移切割patch形成随机点位
        for y in range(0, height, patch_size):
            if y + patch_size > height:
                y = height - patch_size  # 保证patch规格，超出边界往前推
            for x in range(0, width, patch_size):
                if x + patch_size > width:
                    x = width - patch_size
                # 定义块的范围
                y_end = min(y + patch_size, height)
                x_end = min(x + patch_size, width)

                # 切割块
                png_patch = img[y:y_end, x:x_end]
                mask_patch = mask[y:y_end, x:x_end]

                if mask_patch.max() == 0:
                    continue

                png_name = f"{file_name}_{x}_{y}.png"
                txt_name = f"{file_name}_{x}_{y}.txt"

                points = np.argwhere(mask_patch == 1)
                for point in points:
                    label = (
                        0,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')

                Image.fromarray(png_patch).save(png_path + png_name)

                # kernel = np.ones((10, 10), np.uint8)  # 定义膨胀结构元素
                # diffused_matrix = cv2.dilate(mask_patch, kernel, iterations=1)
                # Image.fromarray((diffused_matrix * 255).astype(np.uint8)).save(mask_path + png_name)
                index += 1
